{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Easy Ab initio calculation with ASE-Siesta-Pyscf\n",
    "\n",
    "## No installation necessary, just download a ready to go container for any system, or run it into the cloud"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### We first import the necessary libraries and define the system using ASE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import libraries and set up the molecule geometry\n",
    "\n",
    "from ase.units import Ry, eV, Ha\n",
    "from ase.calculators.siesta import Siesta\n",
    "from ase import Atoms\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from ase.build import molecule\n",
    "\n",
    "CH4 = molecule(\"CH3COOH\")\n",
    "\n",
    "# visualization of the particle\n",
    "from ase.visualize import view\n",
    "view(CH4, viewer='x3d')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### We can then run the DFT calculation using Siesta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# enter siesta input and run siesta\n",
    "siesta = Siesta(\n",
    "    mesh_cutoff=150 * Ry,\n",
    "    basis_set='DZP',\n",
    "    pseudo_qualifier='lda',\n",
    "    energy_shift=(10 * 10**-3) * eV,\n",
    "    fdf_arguments={\n",
    "        'SCFMustConverge': False,\n",
    "        'COOP.Write': True,\n",
    "        'WriteDenchar': True,\n",
    "        'PAO.BasisType': 'split',\n",
    "        'DM.Tolerance': 1e-4,\n",
    "        'DM.MixingWeight': 0.1,\n",
    "        'MaxSCFIterations': 300,\n",
    "        'DM.NumberPulay': 4,\n",
    "        'XML.Write': True})\n",
    "\n",
    "CH4.set_calculator(siesta)\n",
    "e = CH4.get_potential_energy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The TDDFT calculations with PySCF-NAO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute polarizability using pyscf-nao\n",
    "\n",
    "freq = np.arange(0.0, 15.0, 0.05)\n",
    "siesta.pyscf_tddft(label=\"siesta\", jcutoff=7, iter_broadening=0.15/Ha,\n",
    "        xc_code='LDA,PZ', tol_loc=1e-6, tol_biloc=1e-7, freq = freq)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot polarizability with matplotlib\n",
    "%matplotlib inline\n",
    "\n",
    "fig = plt.figure(1)\n",
    "ax1 = fig.add_subplot(121)\n",
    "ax2 = fig.add_subplot(122)\n",
    "ax1.plot(siesta.results[\"freq range\"], siesta.results[\"polarizability nonin\"][:, 0, 0].imag)\n",
    "ax2.plot(siesta.results[\"freq range\"], siesta.results[\"polarizability inter\"][:, 0, 0].imag)\n",
    "\n",
    "ax1.set_xlabel(r\"$\\omega$ (eV)\")\n",
    "ax2.set_xlabel(r\"$\\omega$ (eV)\")\n",
    "\n",
    "ax1.set_ylabel(r\"Im($P_{xx}$) (au)\")\n",
    "ax2.set_ylabel(r\"Im($P_{xx}$) (au)\")\n",
    "\n",
    "ax1.set_title(r\"Non interacting\")\n",
    "ax2.set_title(r\"Interacting\")\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import plotly.offline as py\n",
    "import plotly.graph_objs as go\n",
    "\n",
    "py.init_notebook_mode(connected=True)\n",
    "\n",
    "trace = go.Scatter(x=siesta.results[\"freq range\"], y=siesta.results[\"polarizability inter\"][:, 0, 0].imag)\n",
    "\n",
    "print(CH4.get_positions())\n",
    "py.iplot([trace])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "### Compute the spatial distributoin of the density change at resonance frequency"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "res = 10.5/Ha \n",
    "lim = 20.0 # Bohr\n",
    "box = np.array([[-lim, lim],\n",
    "                [-lim, lim],\n",
    "                [-lim, lim]])\n",
    "from pyscf.nao.m_comp_spatial_distributions import spatial_distribution\n",
    "\n",
    "spd = spatial_distribution(siesta.results[\"density change inter\"], freq/Ha, box, label=\"siesta\")\n",
    "spd.get_spatial_density(8.25/Ha)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "center = np.array([spd.dn_spatial.shape[0]/2, spd.dn_spatial.shape[1]/2, spd.dn_spatial.shape[2]/2], dtype=int)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig2 = plt.figure(2, figsize=(15, 12))\n",
    "\n",
    "cmap=\"seismic\"\n",
    "ax1 = fig2.add_subplot(1, 3, 1)\n",
    "vmax = np.max(abs(spd.dn_spatial[center[0], :, :].imag))\n",
    "vmin = -vmax\n",
    "ax1.imshow(spd.dn_spatial[center[0], :, :].imag, interpolation=\"bicubic\", vmin=vmin, vmax=vmax, cmap=cmap, extent=[spd.mesh[1][0], spd.mesh[1][spd.mesh[1].shape[0]-1], spd.mesh[2][0], spd.mesh[2][spd.mesh[2].shape[0]-1]])\n",
    "\n",
    "ax2 = fig2.add_subplot(1, 3, 2)\n",
    "vmax = np.max(abs(spd.dn_spatial[:, center[1], :].imag))\n",
    "vmin = -vmax\n",
    "ax2.imshow(spd.dn_spatial[:, center[1], :].imag, interpolation=\"bicubic\", vmin=vmin, vmax=vmax, cmap=cmap, extent=[spd.mesh[0][0], spd.mesh[0][spd.mesh[0].shape[0]-1], spd.mesh[2][0], spd.mesh[2][spd.mesh[2].shape[0]-1]])\n",
    "\n",
    "ax3 = fig2.add_subplot(1, 3, 3)\n",
    "vmax = np.max(abs(spd.dn_spatial[:, :, center[2]].imag))\n",
    "vmin = -vmax\n",
    "ax3.imshow(spd.dn_spatial[:, :, center[2]].imag, interpolation=\"bicubic\", vmin=vmin, vmax=vmax, cmap=cmap, extent=[spd.mesh[0][0], spd.mesh[0][spd.mesh[0].shape[0]-1], spd.mesh[1][0], spd.mesh[1][spd.mesh[1].shape[0]-1]])\n",
    "\n",
    "ax1.set_xlabel(r\"y (Bohr)\")\n",
    "ax2.set_xlabel(r\"x (Bohr)\")\n",
    "ax3.set_xlabel(r\"x (Bohr)\")\n",
    "\n",
    "ax1.set_ylabel(r\"z (Bohr)\")\n",
    "ax2.set_ylabel(r\"z (Bohr)\")\n",
    "ax3.set_ylabel(r\"y (Bohr)\")\n",
    "\n",
    "ax1.set_title(r\"Im($\\delta n$) in the $x$ plane\")\n",
    "ax2.set_title(r\"Im($\\delta n$) in the $y$ plane\")\n",
    "ax3.set_title(r\"Im($\\delta n$) in the $z$ plane\")\n",
    "\n",
    "fig2.savefig(\"density_chang_im_CH3COOH.pdf\", format=\"pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
